from datetime import datetime
from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.utils.trigger_rule import TriggerRule
from airflow.exceptions import AirflowSkipException
from airflow.sdk.definitions.edges import EdgeModifier

default_args = {
    'owner': 'airflow',
    'start_date': datetime(2023, 2, 1)
}

def always_skip_function():
    raise AirflowSkipException("This task is intentionally skipped.")

with DAG('1min_file', schedule=None, catchup=True, default_args=default_args):

    # Test case 1: One success, one failure -> Expected state: None
    a_success = EmptyOperator(task_id="a_success")
    a_fails = BashOperator(
        task_id='a_fails',
        bash_command='exit 1',  # A non-zero exit code indicates failure
    )
    rule_1 = EmptyOperator(task_id="rule_1", trigger_rule=TriggerRule.ALL_DONE_MIN_ONE_SUCCESS)

    a_success >> EdgeModifier("expect rule_1 to be successful") >> rule_1
    a_fails >> EdgeModifier("expect rule_1 to be successful") >> rule_1

    # Test case 2: Two successes -> Expected state: None
    b_success = EmptyOperator(task_id="b_success")
    b_success2 = EmptyOperator(task_id="b_success2")
    rule_2 = EmptyOperator(task_id="rule_2", trigger_rule=TriggerRule.ALL_DONE_MIN_ONE_SUCCESS)

    b_success >> EdgeModifier("expect rule_2 to be successful") >>  rule_2
    b_success2 >> EdgeModifier("expect rule_2 to be successful") >> rule_2

    # Test case 3: One success, one skip -> Expected state: SKIPPED
    c_success = EmptyOperator(task_id="c_success")
    c_skip = PythonOperator(
        task_id='c_skip',
        python_callable=always_skip_function,
    )
    rule_3 = EmptyOperator(task_id="rule_3", trigger_rule=TriggerRule.ALL_DONE_MIN_ONE_SUCCESS)

    c_success  >> EdgeModifier("expect rule_3 to be skipped") >> rule_3
    c_skip >> EdgeModifier("expect rule_3 to be skipped") >>  rule_3

    # Test case 4: All failures -> Expected state: UPSTREAM_FAILED
    d_fails1 = BashOperator(
        task_id='d_fails1',
        bash_command='exit 1',
    )
    d_fails2 = BashOperator(
        task_id='d_fails2',
        bash_command='exit 1',
    )
    rule_4 = EmptyOperator(task_id="rule_4", trigger_rule=TriggerRule.ALL_DONE_MIN_ONE_SUCCESS)

    d_fails1 >> EdgeModifier("expect rule_4 to be upstream failed") >> rule_4
    d_fails2 >> EdgeModifier("expect rule_4 to be upstream failed") >> rule_4


    # Test case 5: upstream failure cascade -> Expected state: UPSTREAM_FAILED
    e_fails3 = BashOperator(
        task_id='d_fails3',
        bash_command='exit 1',
    )
    e_upstream_failed = EmptyOperator(task_id="e_upstream")

    rule_5 = EmptyOperator(task_id="rule_5", trigger_rule=TriggerRule.ALL_DONE_MIN_ONE_SUCCESS)

    e_fails3 >> e_upstream_failed >> EdgeModifier("expect rule_5 to be upstream failed") >> rule_5